{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pass2Path Network - Decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import pickle\n",
    "import copy\n",
    "import string\n",
    "import json\n",
    "import csv\n",
    "import time\n",
    "import math\n",
    "from word2keypress import Keyboard\n",
    "from ast import literal_eval\n",
    "from tensorflow.contrib.seq2seq.python.ops import beam_search_ops\n",
    "from pathlib import Path\n",
    "from edit_distance_backtrace import find_med_backtrace\n",
    "import random\n",
    "from collections import OrderedDict\n",
    "from pass2path_model import Pass2PathModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Preprocessing:\n",
    "We need to create look-up tables in order to translate characters to mathematical representation.\n",
    "'''\n",
    "\n",
    "# CODES = {'<PAD>': 0, '<EOS>': 1, '<UNK>': 2, '<GO>': 3 }\n",
    "CODES = {'<GO>': 0, '<EOS>': 1, '<UNK>': 2}\n",
    "# chars = list(string.ascii_letters) + list(string.punctuation) +\n",
    "# list(string.digits) + [\" \", \"\\t\", \"\\x03\", \"\\x04\"]\n",
    "def create_lookup_tables_from_lst(char_lst):\n",
    "    '''\n",
    "    This function creates a dictionary out of a list with the added codes representing padding,\n",
    "    unknows components, start and end of a sequence\n",
    "    '''\n",
    "    # <EOS> acts as <PAD>\n",
    "    # CODES = {'<PAD>': 0, '<EOS>': 1, '<UNK>': 2, '<GO>': 3 }\n",
    "    # chars = list(string.ascii_letters) + list(string.punctuation) +\n",
    "    # list(string.digits) + [\" \", \"\\t\", \"\\x03\", \"\\x04\"]\n",
    "\n",
    "    # make a list of unique chars (from https://stackoverflow.com/a/480227)\n",
    "    seen = set()\n",
    "    seen_add = seen.add\n",
    "    vocab = [x for x in char_lst if not (x in seen or seen_add(x))]\n",
    "\n",
    "    if len(vocab) < 100:\n",
    "        print(vocab)\n",
    "        print(char_lst)\n",
    "    # (1)\n",
    "    # starts with the special tokens\n",
    "    vocab_to_int = copy.copy(CODES)\n",
    "\n",
    "    # the index (v_i) will starts from 4 (the 2nd arg in enumerate() specifies\n",
    "    # the starting index)\n",
    "    # since vocab_to_int already contains special tokens\n",
    "    for v_i, v in enumerate(vocab, len(CODES)):\n",
    "        vocab_to_int[str(v)] = v_i # opposite would be int()\n",
    "    int_to_vocab = {v_i: v for v, v_i in vocab_to_int.items()}\n",
    "\n",
    "    return vocab_to_int, int_to_vocab\n",
    "\n",
    "def build_path_vocab(origin_vocab):\n",
    "    '''\n",
    "    This functions combines the path dictionary with the added codes.\n",
    "    '''\n",
    "    int_lst = sorted(list(origin_vocab.values()))\n",
    "    # print(\"TRANS_t_IDX:{}\".format(int_lst[:10]))\n",
    "    path_to_int_vocab, int_to_path_vocab = create_lookup_tables_from_lst(int_lst)\n",
    "    return path_to_int_vocab, int_to_path_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '!', '\"', '#', '$', '%', '&', \"'\", '(', ')', '*', '+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\\\', ']', '^', '_', '`', '{', '|', '}', '~', ' ', '\\t', '\\x03', '\\x04']\n",
      "['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '!', '\"', '#', '$', '%', '&', \"'\", '(', ')', '*', '+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\\\', ']', '^', '_', '`', '{', '|', '}', '~', ' ', '\\t', '\\x03', '\\x04']\n"
     ]
    }
   ],
   "source": [
    "# Globals:\n",
    "kb = Keyboard()\n",
    "# thisfolder = Path(__file__).absolute().parent\n",
    "# TRANS_to_IDX = json.load((thisfolder / 'data/trans_dict_2idx.json').open()) #\n",
    "# py\n",
    "with open(('./data/trans_dict_2idx.json')) as f:\n",
    "    TRANS_to_IDX = json.load(f) # jupyter\n",
    "IDX_to_TRANS = {v: literal_eval(k) for k, v in TRANS_to_IDX.items()}\n",
    "char_lst = list(string.ascii_letters) + list(string.digits) + list(string.punctuation) + [\" \", \"\\t\", \"\\x03\", \"\\x04\"]\n",
    "source_vocab_to_int, source_int_to_vocab = create_lookup_tables_from_lst(char_lst)\n",
    "target_vocab_to_int, target_int_to_vocab = build_path_vocab(TRANS_to_IDX)\n",
    "latest_ckpt = tf.train.latest_checkpoint('./model/')\n",
    "\n",
    "\n",
    "tf.app.flags.DEFINE_integer('beam_width', 10, 'Beam width used in beamsearch')\n",
    "tf.app.flags.DEFINE_integer('decode_batch_size', 100, 'Batch size used for decoding')\n",
    "tf.app.flags.DEFINE_integer('max_decode_step', 500, 'Maximum time step limit to decode')\n",
    "# tf.app.flags.DEFINE_boolean('write_n_best', False, 'Write n-best list (n=beam_width)')\n",
    "tf.app.flags.DEFINE_string('model_path',latest_ckpt, 'Path to a specific model checkpoint.')\n",
    "\n",
    "tf.app.flags.DEFINE_string('decode_input', 'test_new-opw_5000.txt', 'Decoding input path')\n",
    "tf.app.flags.DEFINE_string('decode_output', 'data/pass2path.predictions', 'Decoding output path')\n",
    "\n",
    "# Runtime parameters\n",
    "tf.app.flags.DEFINE_boolean('allow_soft_placement', True, 'Allow device soft placement')\n",
    "tf.app.flags.DEFINE_boolean('log_device_placement', False, 'Log placement of ops on devices')\n",
    "\n",
    "FLAGS = tf.app.flags.FLAGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_decode_flags(goal='decode', checkpoint=-1, beam_width=10, decode_batch_size=80, tst_file_path='test_new-opw_5000.txt'):\n",
    "    tf.app.flags.FLAGS.__flags.clear()\n",
    "\n",
    "    tf.app.flags.DEFINE_integer('beam_width', beam_width, 'Beam width used in beamsearch')\n",
    "    if (goal == 'decode'):\n",
    "        batch_size = decode_batch_size\n",
    "    else:\n",
    "        batch_size = 1\n",
    "    tf.app.flags.DEFINE_integer('decode_batch_size', batch_size, 'Batch size used for decoding')\n",
    "    tf.app.flags.DEFINE_integer('max_decode_step', 500, 'Maximum time step limit to decode')\n",
    "    # tf.app.flags.DEFINE_boolean('write_n_best', False, 'Write n-best list\n",
    "    # (n=beam_width)')\n",
    "    if (checkpoint == -1):\n",
    "        ckpt = latest_ckpt\n",
    "    else:\n",
    "        ckpt = './model/pass2path.ckpt-' + str(checkpoint)\n",
    "    tf.app.flags.DEFINE_string('model_path',ckpt, 'Path to a specific model checkpoint.')\n",
    "\n",
    "    tf.app.flags.DEFINE_string('decode_input', tst_file_path, 'Decoding input path')\n",
    "    tf.app.flags.DEFINE_string('decode_output', 'data/pass2path_' + str(checkpoint) + '_' + (tst_file_path.split('/'))[-1][:-4] + '.predictions', 'Decoding output path')\n",
    "\n",
    "    # Runtime parameters\n",
    "    tf.app.flags.DEFINE_boolean('allow_soft_placement', True, 'Allow device soft placement')\n",
    "    tf.app.flags.DEFINE_boolean('log_device_placement', False, 'Log placement of ops on devices')\n",
    "\n",
    "    # Ignore Cmmand Line\n",
    "    tf.app.flags.DEFINE_string('x', '', '')\n",
    "    tf.app.flags.DEFINE_string('d', '', '')\n",
    "    tf.app.flags.DEFINE_string('p', '', '')\n",
    "    tf.app.flags.DEFINE_string('o', '', '')\n",
    "    tf.app.flags.DEFINE_string('w', '', '')\n",
    "    tf.app.flags.DEFINE_string('a', '', '')\n",
    "    tf.app.flags.DEFINE_string('b', '', '')\n",
    "    tf.app.flags.DEFINE_string('j', '', '')\n",
    "    tf.app.flags.DEFINE_string('u', '', '')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latest_ckpt = tf.train.latest_checkpoint('./model/')\n",
    "set_decode_flags()\n",
    "FLAGS = tf.app.flags.FLAGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_config(FLAGS):\n",
    "    \n",
    "    config = json.load(open('%s.json' % FLAGS.model_path, 'r'))\n",
    "    for key, value in FLAGS.__flags.items():\n",
    "        config[key] = value.value\n",
    "\n",
    "    return config\n",
    "\n",
    "def load_model(session, config):\n",
    "    \n",
    "    model = Pass2PathModel(config, 'decode')\n",
    "    if tf.train.checkpoint_exists(FLAGS.model_path):\n",
    "        print('Reloading model parameters..')\n",
    "        model.restore(session, FLAGS.model_path)\n",
    "    else:\n",
    "        raise ValueError('No such file:[{}]'.format(FLAGS.model_path))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pass_to_seq(password, vocab_to_int):\n",
    "    '''\n",
    "    This function transforms password to sequence of integers, in order to make the tested password\n",
    "    applicable to the network's input.\n",
    "    '''\n",
    "    results = []\n",
    "    #print(\">>>\", password)\n",
    "    for c in password:\n",
    "        results.append(vocab_to_int.get(c, 2)) # <UNK> is 2\n",
    "    return results\n",
    "\n",
    "\n",
    "def path2word_kb_feasible(word, path, print_path=False):\n",
    "    '''\n",
    "    This function decodes the word in which the given path transitions the input word into.\n",
    "    This is the KeyPress version, which handles the keyboard representations.\n",
    "    If one of the parts components is not feasible (e.g removing a char from out of range index), it skips it\n",
    "    Input parameters: original word, transition path\n",
    "    Output: decoded word\n",
    "    '''\n",
    "    #kb = Keyboard()\n",
    "    word = kb.word_to_keyseq(word)\n",
    "    if not path:\n",
    "        return kb.keyseq_to_word(word)\n",
    "    #path = [literal_eval(p) for p in path]\n",
    "    if (print_path):\n",
    "        print(path)\n",
    "#     print(type(path))\n",
    "#     print(word)\n",
    "    final_word = []\n",
    "    word_len = len(word)\n",
    "    path_len = len(path)\n",
    "    i = 0\n",
    "    j = 0\n",
    "    while (i < word_len or j < path_len):\n",
    "        if ((j < path_len and path[j][2] == i) or (i >= word_len and path[j][2] >= i)):\n",
    "            if (path[j][0] == \"s\"):\n",
    "                # substitute\n",
    "                final_word.append(path[j][1])\n",
    "                i += 1\n",
    "                j += 1\n",
    "            elif (path[j][0] == \"d\"):\n",
    "                # delete\n",
    "                i += 1\n",
    "                j += 1\n",
    "            else:\n",
    "                # \"i\", insert\n",
    "                final_word.append(path[j][1])\n",
    "                j += 1\n",
    "        else:\n",
    "            if (i < word_len):\n",
    "                final_word.append(word[i])\n",
    "                i += 1\n",
    "            if (j < path_len and i > path[j][2]):\n",
    "                j += 1\n",
    "    return (kb.keyseq_to_word(''.join(final_word)))\n",
    "\n",
    "def path_to_pass(password, idx_path, trans_dict):\n",
    "    '''\n",
    "    This function decodes the password in which the given path transitions the input password into.\n",
    "    Input parameters: original password, transition path, transition dictionary\n",
    "    Output: decoded password\n",
    "    '''\n",
    "#     print(idx_path)\n",
    "    str_path = []\n",
    "    for i in idx_path:\n",
    "        if (i != '<PAD>' and i != '<UNK>') and trans_dict.get(int(i)):\n",
    "            str_path.append(trans_dict[int(i)])\n",
    "#         else:\n",
    "#             print(\"could not find \" + str(i) + \" in dictionary\")\n",
    "#     print(str_path)\n",
    "    output_pass = path2word_kb_feasible(password, str_path)\n",
    "    return output_pass\n",
    "\n",
    "def get_accuracy_beam_decode(logits, pass1_batch, pass2_batch, target_int_to_vocab, trans_dict_2path, bias=False):\n",
    "    \"\"\"\n",
    "    Calculate accuracy of BeamSearch output as follows: if one of the K outputs is correct,\n",
    "    it counts as a correct prediction (positive contribution to the accuracy).\n",
    "    \"logits\" are now of the shape: [batches, max_seq, K]\n",
    "    \"\"\"\n",
    "#     print(logits.shape)\n",
    "    beam_width = logits.shape[2]\n",
    "    match_vec = np.zeros((1, logits.shape[0]), dtype=bool)\n",
    "    for i in range(beam_width):\n",
    "        decode_pred_batch = logits[:,:,i]\n",
    "        for k in range(logits.shape[0]):\n",
    "            #print(k)\n",
    "            decode_pred = logits[k,:,i]\n",
    "#             print(decode_pred)\n",
    "            path_idx_pred = [target_int_to_vocab[j] for j in decode_pred]\n",
    "#             print(path_idx_pred)\n",
    "            if ('<EOS>' in path_idx_pred):\n",
    "                path_idx_pred = path_idx_pred[:path_idx_pred.index('<EOS>')]\n",
    "            prediction = path_to_pass(pass1_batch[k], path_idx_pred, trans_dict_2path)\n",
    "            print(\"pred: {}, targ: {}\".format(prediction, pass2_batch[k]))\n",
    "            if (prediction == kb.keyseq_to_word(pass2_batch[k])):\n",
    "#                 if not match_vec[0, k]:\n",
    "#                     print(prediction)\n",
    "                match_vec[0, k] = True\n",
    "            if (bias and (not match_vec[0, k]) and (i == beam_width - 1)):\n",
    "                #print(k)\n",
    "                #print(pass2_batch[k])\n",
    "                #print(kb.keyseq_to_word(pass2_batch[k]))\n",
    "                #print(pass1_batch[k])\n",
    "                #print(kb.keyseq_to_word(pass1_batch[k]))\n",
    "                if (kb.keyseq_to_word(pass2_batch[k]) == kb.keyseq_to_word(pass1_batch[k])):\n",
    "                    match_vec[0, k] = True       \n",
    "    acc = np.mean(match_vec)\n",
    "#     if (max(acc) > 0.6):\n",
    "#         print(target)\n",
    "#         print(np.reshape(np.transpose(logits[:,:,acc.index(max(acc))]),\n",
    "#         target.shape))\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sequence_batch(sequence_batch, pad_int):\n",
    "    \"\"\"\n",
    "    Pad sequences with <PAD> = <EOS> so that each sequence of a batch has the same length\n",
    "    \"\"\"\n",
    "    max_sequence = max([len(seq) for seq in sequence_batch])\n",
    "    return [seq + [pad_int] * (max_sequence - len(seq)) for seq in sequence_batch]\n",
    "def test_samples_gen(fpath):\n",
    "    with open(fpath) as tst_file:\n",
    "        for line in tst_file:\n",
    "            if not all(0 < ord(c) < 255 for c in line): \n",
    "                continue\n",
    "            if ' ' in line:\n",
    "                continue\n",
    "            p = line.rstrip().split('\\t')\n",
    "            if (len(p) != 2):\n",
    "                continue\n",
    "            #if ('\\n' in p[1]):\n",
    "            #    p[1] = p[1][:-1]\n",
    "            #print(line)\n",
    "            #print(p[0])\n",
    "            #print(p[1])\n",
    "            yield p[0], p[1]\n",
    "def test_batches_gen(fpath, batch_size):\n",
    "    samples_gen = test_samples_gen(fpath)\n",
    "    curr_batch_size = 0\n",
    "    source_batch = []\n",
    "    target_batch = []\n",
    "    for i, sample in enumerate(samples_gen):\n",
    "        #print('Sample: ', i)\n",
    "        if (curr_batch_size < batch_size):\n",
    "            source_batch.append(sample[0])\n",
    "            target_batch.append(sample[1])\n",
    "            curr_batch_size += 1\n",
    "        else:\n",
    "            yield source_batch, target_batch\n",
    "            curr_batch_size = 1\n",
    "            source_batch = [sample[0]]\n",
    "            target_batch = [sample[1]]\n",
    "    yield source_batch, target_batch\n",
    "\n",
    "def predict_batches_gen(fpath, batch_size):\n",
    "    curr_batch_size = 0\n",
    "    predict_batch = []\n",
    "    with open(fpath) as in_file:\n",
    "        for line in in_file:\n",
    "            password = line.split('\\n')[0]\n",
    "            if (curr_batch_size < batch_size):\n",
    "                predict_batch.append(password)\n",
    "                curr_batch_size += 1\n",
    "            else:\n",
    "                yield predict_batch\n",
    "                predict_batch = [password]\n",
    "                curr_batch_size = 1\n",
    "        yield predict_batch\n",
    "\n",
    "def preprocess_batch_prediction(batch, source_vocab_to_int):\n",
    "    pass_batch = [kb.word_to_keyseq(d) for d in batch[0]]\n",
    "    pass2_batch = [kb.word_to_keyseq(d) for d in batch[1]]\n",
    "    \n",
    "    batch_lengths = np.array([len(p) for p in pass_batch])\n",
    "    pass_batch_ids = [pass_to_seq(p, source_vocab_to_int) for p in pass_batch]\n",
    "    source_pad_int = source_vocab_to_int['<EOS>']\n",
    "    # Pad\n",
    "    pad_sources_batch = np.array(pad_sequence_batch(pass_batch_ids, source_pad_int))\n",
    "    return (pad_sources_batch, batch_lengths,pass_batch, pass2_batch)\n",
    "\n",
    "def preprocess_batch_prediction_single_pass(batch, source_vocab_to_int):\n",
    "    pass_batch = [kb.word_to_keyseq(d) for d in batch]\n",
    "    batch_lengths = np.array([len(p) for p in pass_batch])\n",
    "    pass_batch_ids = [pass_to_seq(p, source_vocab_to_int) for p in pass_batch]\n",
    "    source_pad_int = source_vocab_to_int['<EOS>']\n",
    "    # Pad\n",
    "    pad_sources_batch = np.array(pad_sequence_batch(pass_batch_ids, source_pad_int))\n",
    "    return (pad_sources_batch, batch_lengths, pass_batch)\n",
    "\n",
    "def create_test_file(csv_path, num_samples=5000, prob=0.001):\n",
    "    '''\n",
    "    This function creates a test file format (\"pass1<tab>pass2\") from a csv dataset\n",
    "    '''\n",
    "    coin = [True, False]\n",
    "    probs = [prob, 1 - prob]\n",
    "    sample_gen = test_samples_csv_gen(csv_path)\n",
    "    count = 0\n",
    "    with open(csv_path[:-4] + '_test.txt', 'w') as tfile:\n",
    "        for i, sample in enumerate(sample_gen):\n",
    "            if (count >= num_samples):\n",
    "                break\n",
    "            if (np.random.choice(coin, 1, p=probs)[0]):\n",
    "                tfile.write(sample[0] + '\\t' + sample[1] +'\\n')\n",
    "                count += 1\n",
    "    print(\"Created test file of {} samples in {}\".format(count, csv_path[:-4] + '_test.txt'))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode(write_to_file=False):\n",
    "    # Load model config\n",
    "    config = load_config(FLAGS)\n",
    "\n",
    "    # Load source data to decode\n",
    "    trans_dict_2path = IDX_to_TRANS\n",
    "    b_gen = test_batches_gen(FLAGS.decode_input, FLAGS.decode_batch_size)\n",
    "    # Initiate TF session\n",
    "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=FLAGS.allow_soft_placement, \n",
    "        log_device_placement=FLAGS.log_device_placement, gpu_options=tf.GPUOptions(allow_growth=True))) as sess:\n",
    "        # Reload existing checkpoint\n",
    "        model = load_model(sess, config)\n",
    "\n",
    "        try:\n",
    "            print('Decoding and Testing {}..'.format(FLAGS.decode_input))\n",
    "            if (write_to_file):\n",
    "                print('Saving predictions in {}'.format(FLAGS.decode_output))\n",
    "                fout = open(FLAGS.decode_output, 'w')\n",
    "            acc = 0\n",
    "            total_samples = 0\n",
    "            for batch_i, batch in enumerate(b_gen):\n",
    "                print(\"Decoding batch # {}\".format(batch_i))\n",
    "                batch_size = len(batch[0])\n",
    "                total_samples += batch_size\n",
    "    #             print(max_seq_len)\n",
    "                pad_encoded_s_batch, batch_lengths, s_batch, t_batch = preprocess_batch_prediction(batch, source_vocab_to_int)\n",
    "                # predicted_ids: GreedyDecoder; [batch_size, max_time_step, 1]\n",
    "                # BeamSearchDecoder; [batch_size, max_time_step, beam_width]\n",
    "                decode_logits, decode_scores = model.predict_scores(sess, encoder_inputs=pad_encoded_s_batch, \n",
    "                                                                    encoder_inputs_length=batch_lengths)\n",
    "                b_acc = get_accuracy_beam_decode(decode_logits, s_batch, t_batch,\n",
    "                                                 target_int_to_vocab, trans_dict_2path, bias=True)\n",
    "                acc += batch_size * b_acc\n",
    "                if (write_to_file):\n",
    "                    # Write decoding results\n",
    "                    for i in range(decode_logits.shape[0]):\n",
    "                        predictions = []\n",
    "                        for k in range(FLAGS.beam_width):\n",
    "                            decode_pred = decode_logits[i,:,k]\n",
    "                #             print(decode_pred)\n",
    "                            path_idx_pred = [target_int_to_vocab[j] for j in decode_pred]\n",
    "                #             print(path_idx_pred)\n",
    "                            if ('<EOS>' in path_idx_pred):\n",
    "                                path_idx_pred = path_idx_pred[:path_idx_pred.index('<EOS>')]\n",
    "                            predictions.append(path_to_pass(s_batch[i], path_idx_pred, trans_dict_2path))\n",
    "                        pred_with_scores = list(zip(predictions, (decode_scores.ravel()).tolist()))\n",
    "                        fout.write(kb.keyseq_to_word(s_batch[i]) + '\\t' + json.dumps(pred_with_scores) + '\\n')\n",
    "                        predictions = []\n",
    "\n",
    "            acc = acc / total_samples\n",
    "#             print(\"Test Accuracy: {0:.4f}\".format(acc))\n",
    "        except IOError:\n",
    "            pass\n",
    "        finally:\n",
    "            if (write_to_file):\n",
    "                fout.close()\n",
    "                print(\"Predictions saved at {}\".format(FLAGS.decode_output))\n",
    "            print(\"Test Accuracy: {0:.4f}\".format(acc))\n",
    "\n",
    "def predict_batch(num_uniuqe_predictions=100, bias=False):\n",
    "    # Load model config\n",
    "    config = load_config(FLAGS)\n",
    "    # Load source data to decode\n",
    "    trans_dict_2path = IDX_to_TRANS\n",
    "    b_gen = predict_batches_gen(FLAGS.decode_input, FLAGS.decode_batch_size)\n",
    "    # Initiate TF session\n",
    "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=FLAGS.allow_soft_placement, \n",
    "        log_device_placement=FLAGS.log_device_placement, gpu_options=tf.GPUOptions(allow_growth=True))) as sess:\n",
    "        # Reload existing checkpoint\n",
    "        model = load_model(sess, config)\n",
    "        try:\n",
    "            print('Decoding and Testing {}..'.format(FLAGS.decode_input))\n",
    "            print('Saving predictions in {}'.format(FLAGS.decode_output))\n",
    "            fout = open(FLAGS.decode_output, 'w')\n",
    "            total_samples = 0\n",
    "            for batch_i, batch in enumerate(b_gen):\n",
    "                print(\"Decoding batch # {}\".format(batch_i))\n",
    "                batch_size = len(batch)\n",
    "                total_samples += batch_size\n",
    "    #             print(max_seq_len)\n",
    "                pad_encoded_s_batch, batch_lengths, s_batch = preprocess_batch_prediction_single_pass(batch, source_vocab_to_int)\n",
    "                # predicted_ids: GreedyDecoder; [batch_size, max_time_step, 1]\n",
    "                # BeamSearchDecoder; [batch_size, max_time_step, beam_width]\n",
    "                decode_logits, decode_scores = model.predict_scores(sess, encoder_inputs=pad_encoded_s_batch, \n",
    "                                                                    encoder_inputs_length=batch_lengths)\n",
    "                # Write decoding results\n",
    "                for i in range(decode_logits.shape[0]):\n",
    "                    predictions = []\n",
    "                    for k in range(FLAGS.beam_width):\n",
    "                        decode_pred = decode_logits[i,:,k]\n",
    "                        path_idx_pred = [target_int_to_vocab[j] for j in decode_pred]\n",
    "                        if ('<EOS>' in path_idx_pred):\n",
    "                            path_idx_pred = path_idx_pred[:path_idx_pred.index('<EOS>')]\n",
    "                        predictions.append(path_to_pass(s_batch[i], path_idx_pred, trans_dict_2path))\n",
    "                    #pred_with_scores = list(zip(predictions,\n",
    "                    #(decode_scores.ravel()).tolist()))\n",
    "                    # Take uniuqe predicions and add the original password as a\n",
    "                    # guess\n",
    "                    orig_pass = kb.keyseq_to_word(s_batch[i])\n",
    "                    seen = set()\n",
    "                    seen_add = seen.add\n",
    "                    if (bias):\n",
    "                        unq_predictions = [orig_pass]\n",
    "                        seen_add(orig_pass)\n",
    "                    else:\n",
    "                        unq_predictions = []\n",
    "                    unq_predictions += [x for x in predictions if not (x in seen or seen_add(x))]\n",
    "                    unq_predictions = unq_predictions[:num_uniuqe_predictions]\n",
    "                    fout.write(orig_pass + '\\t' + json.dumps(unq_predictions) + '\\n')\n",
    "                    predictions = []\n",
    "        except IOError:\n",
    "            pass\n",
    "        finally:\n",
    "            fout.close()\n",
    "            print(\"Predictions saved at {}\".format(FLAGS.decode_output))\n",
    "\n",
    "def predict(original_pass):\n",
    "    # Load model config\n",
    "    config = load_config(FLAGS)\n",
    "    trans_dict_2path = IDX_to_TRANS\n",
    "    \n",
    "    decoded_origin_pass = pass_to_seq(original_pass, source_vocab_to_int)\n",
    "    # Initiate TF session\n",
    "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=FLAGS.allow_soft_placement, \n",
    "            log_device_placement=FLAGS.log_device_placement,\n",
    "            gpu_options=tf.GPUOptions(allow_growth=True))) as sess:\n",
    "\n",
    "        # Reload existing checkpoint\n",
    "        model = load_model(sess, config)\n",
    "        print('Decoding...', decoded_origin_pass)\n",
    "        decode_logits = model.predict(sess, encoder_inputs=np.array([decoded_origin_pass]), \n",
    "                                      encoder_inputs_length=np.array([len(decoded_origin_pass)]))\n",
    "        for i in range(decode_logits.shape[0]):\n",
    "            predictions = []\n",
    "            for k in range(FLAGS.beam_width):\n",
    "                decode_pred = decode_logits[i, :, k]\n",
    "                path_idx_pred = [target_int_to_vocab[j] for j in decode_pred]\n",
    "                if ('<EOS>' in path_idx_pred):\n",
    "                    path_idx_pred = path_idx_pred[:path_idx_pred.index('<EOS>')]\n",
    "                predictions.append(path_to_pass(original_pass, path_idx_pred, trans_dict_2path))\n",
    "        print(\"Predictions for {}: {}\".format(original_pass, ', '.join(sorted(predictions))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def predict(original_pass):\n",
    "#     # Load model config\n",
    "#     config = load_config(FLAGS)\n",
    "#     trans_dict_2path = IDX_to_TRANS\n",
    "#     decoded_origin_pass = pass_to_seq(original_pass, source_vocab_to_int)\n",
    "#     # Initiate TF session\n",
    "#     with tf.Session(config=tf.ConfigProto(allow_soft_placement=FLAGS.allow_soft_placement, \n",
    "#         log_device_placement=FLAGS.log_device_placement, gpu_options=tf.GPUOptions(allow_growth=True))) as sess:\n",
    "\n",
    "#         # Reload existing checkpoint\n",
    "#         model = load_model(sess, config)\n",
    "#         print('Decoding...')\n",
    "#         decode_logits = model.predict(sess, encoder_inputs=np.array([decoded_origin_pass]), \n",
    "#                                               encoder_inputs_length=np.array([len(decoded_origin_pass)]))\n",
    "#         for i in range(decode_logits.shape[0]):\n",
    "#             predictions = []\n",
    "#             for k in range(FLAGS.beam_width):\n",
    "#                 decode_pred = decode_logits[i,:,k]\n",
    "#                 path_idx_pred = [target_int_to_vocab[j] for j in decode_pred]\n",
    "#                 if ('<EOS>' in path_idx_pred):\n",
    "#                     path_idx_pred = path_idx_pred[:path_idx_pred.index('<EOS>')]\n",
    "#                 predictions.append(path_to_pass(original_pass, path_idx_pred, trans_dict_2path))\n",
    "#         print(\"Predictions for {}: {}\".format(original_pass, ', '.join(predictions)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "building model..\n",
      "building encoder..\n",
      "building decoder...\n",
      "use beamsearch decoding..\n",
      "building beamsearch decoder..\n",
      "Reloading model parameters..\n",
      "INFO:tensorflow:Restoring parameters from ./model/pass2path.ckpt-1702000\n",
      "model restored from ./model/pass2path.ckpt-1702000\n",
      "Decoding and Testing test_new-opw_5000.txt..\n",
      "Decoding batch # 0\n",
      "Decoding batch # 1\n",
      "Decoding batch # 2\n",
      "Decoding batch # 3\n",
      "Decoding batch # 4\n",
      "Decoding batch # 5\n",
      "Decoding batch # 6\n",
      "Decoding batch # 7\n",
      "Decoding batch # 8\n",
      "Decoding batch # 9\n",
      "Decoding batch # 10\n",
      "Decoding batch # 11\n",
      "Decoding batch # 12\n",
      "Decoding batch # 13\n",
      "Decoding batch # 14\n",
      "Decoding batch # 15\n",
      "Decoding batch # 16\n",
      "Decoding batch # 17\n",
      "Decoding batch # 18\n",
      "Decoding batch # 19\n",
      "Decoding batch # 20\n",
      "Decoding batch # 21\n",
      "Decoding batch # 22\n",
      "Decoding batch # 23\n",
      "Decoding batch # 24\n",
      "Decoding batch # 25\n",
      "Decoding batch # 26\n",
      "Decoding batch # 27\n",
      "Decoding batch # 28\n",
      "Decoding batch # 29\n",
      "Decoding batch # 30\n",
      "Decoding batch # 31\n",
      "Decoding batch # 32\n",
      "Decoding batch # 33\n",
      "Decoding batch # 34\n",
      "Decoding batch # 35\n",
      "Decoding batch # 36\n",
      "Decoding batch # 37\n",
      "Decoding batch # 38\n",
      "Decoding batch # 39\n",
      "Decoding batch # 40\n",
      "Decoding batch # 41\n",
      "Decoding batch # 42\n",
      "Decoding batch # 43\n",
      "Decoding batch # 44\n",
      "Decoding batch # 45\n",
      "Decoding batch # 46\n",
      "Decoding batch # 47\n",
      "Decoding batch # 48\n",
      "Decoding batch # 49\n",
      "Test Accuracy: 0.0251\n"
     ]
    },
    {
     "ename": "SystemExit",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "An exception has occurred, use %tb to see the full traceback.\n",
      "\u001b[1;31mSystemExit\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\ProgramData\\Anaconda3\\envs\\tensorflow\\lib\\site-packages\\IPython\\core\\interactiveshell.py:2971: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
      "  warn(\"To exit: use 'exit', 'quit', or Ctrl-D.\", stacklevel=1)\n"
     ]
    }
   ],
   "source": [
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\" # so the IDs match nvidia-smi\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # \"0, 1\" for multiple\n",
    "# devices = [\"/gpu:0\"]\n",
    "# os.environ['PYTHONHASHSEED'] = '0'\n",
    "# os.environ['TF_CUDNN_USE_AUTOTUNE'] = '0'\n",
    "def main(_):\n",
    "    decode(write_to_file=False)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    tf.reset_default_graph()\n",
    "#     random.seed(0)\n",
    "#     np.random.seed(0)\n",
    "#     tf.set_random_seed(0)\n",
    "    tf.app.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "decode()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csv_path = 'cleaned_pw_paths_tr_10000.csv'\n",
    "create_test_file(csv_path, num_samples=10, prob=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coin = [True, False]\n",
    "probs = [0.001, 1 - 0.001]\n",
    "np.random.choice(coin, 1, p=probs)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./model/pass2path.ckpt-1702000\n"
     ]
    }
   ],
   "source": [
    "print(latest_ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_samples_csv_gen(csv_path):\n",
    "    '''\n",
    "    This function generates pairs of passwods from a given csv file,\n",
    "    assuming the first 2 columns are words\n",
    "    '''\n",
    "    with open(csv_path) as csv_file:\n",
    "        csv_reader = csv.reader(csv_file, delimiter=',')\n",
    "        for row in csv_reader:\n",
    "            yield row[0], row[1]\n",
    "\n",
    "def create_test_file(csv_path, num_samples=5000, prob=0.001):\n",
    "    '''\n",
    "    This function creates a test file format (\"pass1<tab>pass2\") from a csv dataset\n",
    "    '''\n",
    "    coin = [True, False]\n",
    "    probs = [prob, 1 - prob]\n",
    "    sample_gen = test_samples_csv_gen(csv_path)\n",
    "    count = 0\n",
    "    with open(csv_path[:-4] + '_test.txt', 'w') as tfile:\n",
    "        for i, sample in enumerate(sample_gen):\n",
    "            if (count >= num_samples):\n",
    "                break\n",
    "            if (np.random.choice(coin, 1, p=probs)[0]):\n",
    "                tfile.write(sample[0] + '\\t' + sample[1] + '\\n')\n",
    "                count += 1\n",
    "    print(\"Created test file of {} samples in {}\".format(count, csv_path[:-4] + '_test.txt'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
